"""
演示示例：代码补全
使用方法：在代码最后一行，回车，等待 AI 显示出代码补全建议后，按 Tab  键接受
"""

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

# 加载数据
data = pd.read_csv('house_prices.csv')

# 数据预处理
data['Age'] = 2023 - data['YearBuilt']
data = data.dropna()

# 特征选择
features = ['SquareFeet', 'Bedrooms', 'Bathrooms', 'Age']
target = 'Price'

X = data[features]
y = data[target]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建和训练模型
model = LinearRegression()
model.fit(X_train, y_train)

# 预测
y_pred = model.predict(X_test)

# 评估模型
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"均方误差: {mse}")
print(f"R2 分数: {r2}")

# 可视化结果
plt.figure(figsize=(10, 6))
sns.scatterplot(x=y_test, y=y_pred)
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
plt.xlabel('实际价格')
plt.ylabel('预测价格')
plt.title('实际价格 vs 预测价格')
plt.show()

# 特征重要性
importance = pd.DataFrame({'feature': features, 'importance': model.coef_})
importance = importance.sort_values('importance', ascending=False)
plt.figure(figsize=(10, 6))